#include "minimum_snap_path_plan/minimum_snap_path_plan.h"


std::vector<double> getMinimum_X_PlanTrajTotalPointNum(MatrixXd point,double rate,std::vector<double> start_end_va ,double set_speed,double set_acc, 
													int option, int& traj_total_point_num )
{
	double TIME_INTERVAL=1.0/rate;
	int point_num=point.rows();
	const int piece_num=point_num-1;

	double traj_total_time=0;
	traj_total_point_num=0;

	std::vector<double> T;
	T.push_back(0);
    for(int i=0;i<piece_num;i++)
    {
		double delt_x=point(i+1,0)-point(i,0);
		double delt_y=point(i+1,1)-point(i,1);
		double delt_z=point(i+1,2)-point(i,2);

        double t; 
		double T1,T2,T3;

		double v_max=set_speed;
		double a_max=set_acc;

		if(fabs(delt_x)<=pow(v_max,2)/a_max)
		{
			T1=sqrt(4*fabs(delt_x)/a_max);
		}
		else
			T1=(fabs(delt_x)+pow(v_max,2)/a_max)/v_max;
		if(fabs(delt_y)<=pow(v_max,2)/a_max)
		{
			T2=sqrt(4*fabs(delt_y)/a_max);
		}
		else
			T2=(fabs(delt_y)+pow(v_max,2)/a_max)/v_max;

		if(fabs(delt_z)<=pow(v_max,2)/a_max)
		{
			T3=sqrt(4*fabs(delt_z)/a_max);
		}
		else
			T3=(fabs(delt_z)+pow(v_max,2)/a_max)/v_max;

		// T1=(point.row(i+1)-point.row(i)).norm()/v_max;
		// T2=0;T3=0;
		
		// ROS_INFO("T1 T2 T3 %f %f %f",T1,T2,T3);
		double T_initial = T1 > T2 ? T1 : T2;
		t = T_initial> T3 ? T_initial : T3;

		int traj_point_num=t/TIME_INTERVAL;
		traj_point_num=std::max(1,traj_point_num);
		t=traj_point_num*TIME_INTERVAL;
		traj_total_point_num+=traj_point_num;

        T.push_back(t);
		traj_total_time+=t;
    } 
	traj_total_point_num++;


	return T;

}

std::vector<Eigen::Vector3d> Minimum_X_Plan(MatrixXd point,double rate,std::vector<double> start_end_va ,double set_speed,double set_acc,int option,
											int setTotalTrajPointNum)
{
	//option=0:Minimum_Snap, option=1:Minimum_Acceleration


	int order_set[]={5,3};  
	const int order=order_set[option];  // order of the trajectory
	const int n_coefficient=order+1; 	// num of coefficient in a Polynomial trajectory

	int COST_FUNCTION=option;
	// COST_FUNCTION=0;   //Minimum_snap; specify start and end point's position and velocity and acceleration, specify waypoint position
	// COST_FUNCTION=1;   //Minimum_acc;  specify start and end point's position and velocity, specify waypoint position

    double TIME_INTERVAL=1.0/rate;

    // ros::Time start=ros::Time::now();

    int point_num=point.rows();
    const int piece_num=point_num-1;

    // ROS_INFO("num of waypoint in initial path including start and end point   %d",point_num);

    std::vector<double> T;
	int traj_total_point_num=0;
	T=getMinimum_X_PlanTrajTotalPointNum(point, rate, start_end_va , set_speed, set_acc,  option,  traj_total_point_num );
	double originTraj_Total_Time=std::accumulate(T.begin(),T.end(),0.0);

	double setTotalTrajTime=(setTotalTrajPointNum-1)*TIME_INTERVAL;
	double time_ratio=setTotalTrajTime/originTraj_Total_Time;
	for(int i=0;i<(int)T.size();i++)
		T[i]*=time_ratio;

	// ROS_INFO("traj total time: %f, traj point total num: %d, time interval: %f",setTotalTrajTime,traj_total_point_num,TIME_INTERVAL);

	//  traj_total_point_num=total_time/TIME_INTERVAL;
	// TIME_INTERVAL=total_time/(double)traj_total_point_num;

    MatrixXd M=MatrixXd::Zero(n_coefficient*piece_num,n_coefficient*piece_num);
    MatrixXd Q=MatrixXd::Zero(n_coefficient*piece_num,n_coefficient*piece_num);

    for(int i=0;i<piece_num;i++)
    {
        double t=T[i+1];
        MatrixXd m=MatrixXd::Zero(n_coefficient,n_coefficient);
		if(option==0)
		{
			m<< 0,0,0,0,0,1,
				0,0,0,0,1,0,
				0,0,0,2,0,0,
				pow(t,5),pow(t,4),pow(t,3),pow(t,2),pow(t,1),1,
				5*pow(t,4),4*pow(t,3),3*pow(t,2),2*pow(t,1),1,0,
				20*pow(t,3),12*pow(t,2),6*pow(t,1),2,0,0;
		}
		if(option==1)
		{
			m<< 0,0,0,1,
				0,0,1,0,
				pow(t,3),pow(t,2),pow(t,1),1,
				3*pow(t,2),2*pow(t,1),1,0;
		}


        M.block(i*n_coefficient,i*n_coefficient,n_coefficient,n_coefficient)=m;

        MatrixXd q = MatrixXd::Zero(n_coefficient,n_coefficient);

		if(COST_FUNCTION==0)
		{
			Eigen::MatrixXd tmp=MatrixXd::Zero(2,2);
			tmp<<4800*pow(t,3),1440*pow(t,2),
				1440*pow(t,2),576*pow(t,1);
			q.block(0,0,tmp.rows(),tmp.cols())=tmp;
		}
		if(COST_FUNCTION==1)
		{
			Eigen::MatrixXd tmp=MatrixXd::Zero(2,2);
			tmp<<12*pow(t,3),6*pow(t,2),
				6*pow(t,2),4*pow(t,1);
			q.block(0,0,tmp.rows(),tmp.cols())=tmp;			
		}

        // cout<<"before q: \n"<<q<<endl;
		// q = MatrixXd::Zero(n_coefficient,n_coefficient);
        // for(int j=5;j>=4;j--)
        // {
        //     for(int l=5;l>=4;l--)
        //     {
        //         q(5-j,5-l)=j*(j-1)*(j-2)*(j-3)*l*(l-1)*(l-2)*(l-3)/(j+l-7)*pow(t,j+l-7);
        //     }
        // }
        // cout<<"q: ////////////////\n"<<q<<endl;

        Q.block(i*n_coefficient,i*n_coefficient,n_coefficient,n_coefficient)=q;
    }

    // cout<<"M: "<<endl<<M<<endl;
    // cout<<"Q: \n"<<endl<<Q<<endl;

    MatrixXd C_t ;

    int d_f_number=0;
    int d_p_number=0;

	if(COST_FUNCTION==0)
	{
		C_t=MatrixXd::Zero(piece_num*n_coefficient, 3*(point_num)); 
		d_f_number=piece_num+5;
		d_p_number=3*point_num-d_f_number;
		C_t.block<4,4>(0,0)=Matrix<double,4,4>::Identity();

		for(int i=0;i<piece_num-1;i++)
		{
			C_t(i*6+3,i+3)=1;
			C_t(i*6+4,d_f_number+i*2)=1;
			C_t(i*6+5,d_f_number+i*2+1)=1;

			C_t(i*6+6,i+3)=1;
			C_t(i*6+7,d_f_number+i*2)=1;
			C_t(i*6+8,d_f_number+i*2+1)=1;

		}
	    C_t.block<3,3>(6*piece_num-3,2+piece_num)=Matrix<double,3,3>::Identity();
		// cout<<"C_t: "<<endl<<C_t<<endl;
	}

	if(COST_FUNCTION==1)
	{
		C_t=MatrixXd::Zero(4*point_num-4, 2*(point_num)); 

		d_f_number=point_num+2;
		d_p_number=2*point_num-d_f_number;
		C_t.block<2,2>(0,0)=Matrix<double,2,2>::Identity();

		for(int i=0;i<point_num-2;i++)
		{
			C_t(i*4+2,i+2)=1;
			C_t(i*4+3,d_f_number+i)=1;

			C_t(i*4+4,i+2)=1;
			C_t(i*4+5,d_f_number+i)=1;

		}

		C_t(4*point_num-6,point_num)=1;
		C_t(4*point_num-5,point_num+1)=1;


	}

	// std::cout<<"C_t: \n"<<C_t<<std::endl;

    MatrixXd C=C_t.transpose();

	Eigen::MatrixXd M_inverse=M.inverse();

    MatrixXd R=C*M_inverse.transpose()*Q*M_inverse*C_t;

    // cout<<"R: "<<endl<<R<<endl;

    MatrixXd R_pp=R.bottomRightCorner(d_p_number,d_p_number);
    MatrixXd R_fp=R.topRightCorner(d_f_number,d_p_number);

    MatrixXd d_f=MatrixXd::Zero(d_f_number, 3);

	if(COST_FUNCTION==0)
	{
		d_f.row(0)=point.row(0);
		d_f.row(1)<<start_end_va[0],start_end_va[1],0;   //initial velocity constraint
		d_f.row(2)<<start_end_va[2],start_end_va[3],0;   //initial accel constraint
		for(int i=0;i<point_num-1;i++)
    	{
        	d_f.row(3+i)=point.row(i+1);
    	}

    	d_f.row(d_f_number-2)<<start_end_va[4],start_end_va[5],0;  //end velocity constraint
    	d_f.row(d_f_number-1)<<start_end_va[6],start_end_va[7],0; //end accel constraint
	}

	if(COST_FUNCTION==1)
	{
		d_f.row(0)=point.row(0);
		d_f.row(1)<<start_end_va[0],start_end_va[1],0;   //initial velocity constraint
		for(int i=0;i<point_num-1;i++)
    	{
        	d_f.row(2+i)=point.row(i+1);
    	}
    	d_f.row(d_f_number-1)<<start_end_va[4],start_end_va[5],0; //end velocity constraint
	}



    MatrixXd trans_matrix=-R_pp.inverse()*R_fp.transpose();

    std::vector<std::vector<double>> plan_point_plot(3);
    std::vector<std::vector<double>> plan_vel_plot(3);
    std::vector<std::vector<double>> plan_acce_plot(3);

	// ROS_INFO("time before loop: %f",(ros::Time::now()-start).toSec());

	std::vector<Eigen::Vector3d> plan_traj;
	std::vector<Eigen::Vector3d> plan_traj_vel;
	std::vector<Eigen::Vector3d> plan_traj_acc;

	Eigen::MatrixXd p=Eigen::MatrixXd::Zero(n_coefficient*piece_num,3);


	// std::cout<<"M_inverse: \n"<<M_inverse<<std::endl;

	for(int j=0;j<3;j++)
	{
		VectorXd dp_star=trans_matrix*d_f.col(j);
		VectorXd d(d_f_number+d_p_number);
		d.head(d_f_number)=d_f.col(j);
		d.tail(d_p_number)=dp_star;
		
		p.col(j)=M_inverse*C_t*d;		
	}

	// std::cout<<"p: \n"<<p<<std::endl;
	int piece_id=0;
	double time=0;
	int point_count=0;

	while(1)
	{
		if(time>T[piece_id+1] && piece_id!=point_num-2 )
		{
			time=time-T[piece_id+1];
			piece_id++;
		}
		Eigen::Vector3d plan_traj_point;
		Eigen::Vector3d plan_traj_point_vel;
		Eigen::Vector3d plan_traj_point_acc;


		for(int j=0;j<3;j++)
		{
			double plan_point=0,plan_point_vel=0,plan_point_acc=0;
			for(int k=order;k>=0;k--)
			{
				plan_point+=p(order-k+piece_id*n_coefficient,j)*pow(time,k);
				if(k<1)continue;
				plan_point_vel+=k*p(order-k+piece_id*n_coefficient,j)*pow(time,k-1);
				if(k<2)continue;
				plan_point_acc+=k*(k-1)*p(order-k+piece_id*n_coefficient,j)*pow(time,k-2);
			}
			plan_traj_point(j)=plan_point;
			plan_traj_point_vel(j)=plan_point_vel;
			plan_traj_point_acc(j)=plan_point_acc;
			plan_point_plot[j].push_back(plan_point);

		}
		// ROS_INFO("plan_traj_point: %f %f %f",plan_traj_point(0),plan_traj_point(1),plan_traj_point(2));
		plan_traj.push_back(plan_traj_point);
		plan_traj_vel.push_back(plan_traj_point_vel);
		plan_traj_acc.push_back(plan_traj_point_acc);
		time+=TIME_INTERVAL;
		point_count++;

		// if(piece_id==point_num-2 &&  time>T[piece_id+1])
		// 	break;

		if((int)plan_traj.size()==setTotalTrajPointNum)break;
		
	}

	// for(int i=0;i<plan_traj.size();i++)
	// {
	// 	ROS_INFO("polynomial traj point: %f %f %f",plan_traj[i](0),plan_traj[i](1),plan_traj[i](2));
	// }
    // ros::Time end=ros::Time::now();
    // ROS_INFO("calculate polynomial trajectory spend %f traj total point num: %ld traj total time: %f",(end-start).toSec(),plan_traj.size(),setTotalTrajTime);

	// std::ofstream outFile;
	// outFile.open("/home/pp/catkin_ws/src/trajectory_formation/traj.csv"); // 打开模式可省略
	// for(int i=0;i<plan_traj.size();i++)
    // {
    //     outFile << plan_traj[i](0) << ',' << plan_traj[i](1) << ',' << plan_traj[i](2) <<','
    //      << plan_traj_vel[i](0)<<',' <<plan_traj_vel[i](1)<<',' << plan_traj_vel[i](2)
    //             <<',' << plan_traj_acc[i](0)<<',' << plan_traj_acc[i](1)<<',' << plan_traj_acc[i](2)<<std::endl;
    // }
	// outFile.close();
    
    // std::vector<std::vector<double>> origin_point(3);
    // for(int i=0;i<point_num;i++)
    // {
    //     origin_point[0].push_back(point(i,0));
    //     origin_point[1].push_back(point(i,1));
    //     origin_point[2].push_back(point(i,2));
    // }
    // plt::plot(origin_point[0],origin_point[1]);
    // plt::scatter(plan_point_plot[0], plan_point_plot[1]);
    // plt::grid(true);
	// plt::xlim(0, 13);
	// plt::ylim(0, 13);
	// plt::axis("equal");
    // plt::show();

    return plan_traj;
}





std::vector<Eigen::Vector3d> Navie_Plan(MatrixXd point,double rate,std::vector<double> start_end_va ,double set_speed,double set_acc,int option,
											int setTotalTrajPointNum)
{

    double TIME_INTERVAL=1.0/rate;
    int point_num=point.rows();
    const int piece_num=point_num-1;

    std::vector<double> T;
	int traj_total_point_num=0;
	T=getMinimum_X_PlanTrajTotalPointNum(point, rate, start_end_va , set_speed, set_acc,  option,  traj_total_point_num );
	double originTraj_Total_Time=std::accumulate(T.begin(),T.end(),0.0);

	double setTotalTrajTime=(setTotalTrajPointNum-1)*TIME_INTERVAL;
	double time_ratio=setTotalTrajTime/originTraj_Total_Time;
	for(int i=0;i<(int)T.size();i++)
		T[i]*=time_ratio;


    std::vector<std::vector<double>> plan_point_plot(3);
    std::vector<std::vector<double>> plan_vel_plot(3);
    std::vector<std::vector<double>> plan_acce_plot(3);


	std::vector<Eigen::Vector3d> plan_traj;
	std::vector<Eigen::Vector3d> plan_traj_vel;
	std::vector<Eigen::Vector3d> plan_traj_acc;

	int piece_id=0;
	double time=0;
	int point_count=0;

	while(1)
	{
		if(time>T[piece_id+1] && piece_id!=point_num-2 )
		{
			time=time-T[piece_id+1];
			piece_id++;
		}
		Eigen::Vector3d plan_traj_point;
		Eigen::Vector3d plan_traj_point_vel;
		Eigen::Vector3d plan_traj_point_acc;


		Eigen::Vector3d piece_start_point=point.row(piece_id);
		Eigen::Vector3d piece_end_point=point.row(piece_id+1);

		Eigen::Vector3d v_direction=(piece_end_point-piece_start_point).normalized();
		set_speed=(piece_end_point-piece_start_point).norm()/T[piece_id+1];
		Eigen::Vector3d pos=piece_start_point+time*set_speed*v_direction;
		

		plan_traj.push_back(pos);
		plan_traj_vel.push_back(set_speed*v_direction);
		plan_point_plot[0].push_back(pos(0));
		plan_point_plot[1].push_back(pos(1));


		time+=TIME_INTERVAL;
		point_count++;


		if((int)plan_traj.size()==setTotalTrajPointNum)break;
		
	}

	// for(int i=0;i<plan_traj.size();i++)
	// {
	// 	ROS_INFO("polynomial traj point: %f %f %f",plan_traj[i](0),plan_traj[i](1),plan_traj[i](2));
	// }
    // ros::Time end=ros::Time::now();
    // ROS_INFO(" traj total point num: %ld traj total time: %f",plan_traj.size(),setTotalTrajTime);

	// std::ofstream outFile;
	// outFile.open("/home/pp/catkin_ws/src/path_planning_test/traj.csv"); // 打开模式可省略
	// for(int i=0;i<plan_traj.size();i++)
    // {
    //     outFile << plan_traj[i](0) << ',' << plan_traj[i](1) << ',' << plan_traj[i](2)<<","
	// 	 << plan_traj_vel[i](0) << ',' << plan_traj_vel[i](1) << ',' << plan_traj_vel[i](2)  <<std::endl;
    // }
	// outFile.close();
    
    // std::vector<std::vector<double>> origin_point(3);
    // for(int i=0;i<point_num;i++)
    // {
    //     origin_point[0].push_back(point(i,0));
    //     origin_point[1].push_back(point(i,1));
    //     origin_point[2].push_back(point(i,2));
    // }
    // plt::plot(origin_point[0],origin_point[1]);
    // plt::scatter(plan_point_plot[0], plan_point_plot[1]);
    // plt::grid(true);
	// plt::xlim(0, 13);
	// plt::ylim(0, 13);
	// plt::axis("equal");
    // plt::show();

    return plan_traj;
}